-
Notifications
You must be signed in to change notification settings - Fork 136
Accurate Tensor.device
for TFEager backends
#1077
Conversation
// Parse type and ordinal from a string with the expected syntax: | ||
// /job:localhost/replica:0/task:0/device:CPU:0 | ||
let pattern = ".+device:(.+):(\\d+)$" | ||
let regex = try! NSRegularExpression(pattern: pattern) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This string parsing looks expensive. Might want to double check that this doesn't happen too much.
Device.defaultTFEager
Tensor.device
for TFEager backends
|
||
// Parse type and ordinal from a string with the expected syntax: | ||
// /job:localhost/replica:0/task:0/device:CPU:0 | ||
let pattern = ".+device:(.+):(\\d+)$" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe break this String -> Device out as a separate function?
Also, I'm concerned that the string parsing will be expensive. This function is called a lot. (whenever there is a scalar constant or anything like that). I think it would be best to see if you can add your own TFE_TensorHandleDevice_Type and TFE_TensorHandleDevice_Id. Some benchmarking results might work instead.
Different approach coming. Closing this PR |
This modifies
Tensor.device
for TFEager backends to retrieve device details from the underlying eager libraries. This is more accurate than returningDevice.defaultTFEager
and reflects the actual location of eager operations.This was tested by building a CUDA toolchain and verifying that
TFE_TensorHandleDeviceName()
returns:and a GPU device is returned by
Tensor.device
:On a device without GPU,
TFE_TensorHandleDeviceName()
returns:and a CPU device is returned by
Tensor.device
:Fixes tensorflow/swift#524.